import tkinter as tk
from tkinter import filedialog
from tkinter import ttk
import shutil
import os
import hashlib
import mysql.connector
from cryptography.fernet import Fernet
import binascii
import base64
import tkinter.messagebox as messagebox


def checkUser():
    user_id = user_id_entry.get()
    user_pwd = user_id_pw.get()

    query = "SELECT userid FROM user_info WHERE userid = %s AND userpwd = %s"
    cursor.execute(query, (user_id, user_pwd))
    result = cursor.fetchone()
    if cursor.fetchone():
        print("checkUser:having selected more than one user!")
        cursor.fetchall()

    global isUserChecked
    if result:
        isUserChecked = True
        status_label1.config(text="登录成功！")
    else:
        isUserChecked = False
        status_label2.config(text="登录失败")


def calculate_hash(file_path):
    hasher = hashlib.sha256()
    with open(file_path, 'rb') as f:
        while True:
            data = f.read(65536)
            if not data:
                break
            hasher.update(data)
    return hasher.hexdigest()


def uploadFile():
    if not isUserChecked:
        status_label2.config(text="请先登录！")
        return

    user_id = user_id_entry.get()

    file_path = filedialog.askopenfilename()
    if file_path:
        hash_value = calculate_hash(file_path)

        query = "SELECT ownerid, findex FROM file_info WHERE fhash = %s"
        cursor.execute(query, (hash_value,))
        # existing_file = cursor.fetchone()
        # if cursor.fetchone():
        #     print("uploadFile having got more file!")
        #     cursor.fetchall()
        same_id = False
        existing_files = cursor.fetchall()
        for existing_file in existing_files:
            existing_owner_id, existing_file_path = existing_file
            if existing_owner_id == user_id:
                same_id = True
                break

        # 存在相同哈希值文件
        if existing_files:
            existing_owner_id, existing_file_path = existing_files[0]
            if same_id:
                # 生成并回显完整校验码
                encrypted_hash = calculate_hash(existing_file_path)
                windowShowHash(encrypted_hash)
                status_label2.config(text="文件已存在！")
                return
            else:
                insert_query = "INSERT INTO file_info (ownerid, findex, fhash) VALUES (%s, %s, %s)"
                cursor.execute(insert_query, (user_id, existing_file_path, hash_value))
                conn.commit()

                # 生成并回显完整校验码
                encrypted_hash = calculate_hash(existing_file_path)
                windowShowHash(encrypted_hash)
                status_label2.config(text="文件已上传！")

        # 不存在相同哈希值文件，则对目标文件先加密再转存到存储区下
        else:
            destination_dir = os.path.join(r"E:\A课程资料\应用密码学\EP2", user_id)
            os.makedirs(destination_dir, exist_ok=True)

            key_bytes = binascii.unhexlify(hash_value)
            key_base64 = base64.urlsafe_b64encode(key_bytes)
            cipher_suite = Fernet(key_base64)

            with open(file_path, "rb") as file:
                data = file.read()
            encrypted_data = cipher_suite.encrypt(data)
            # with open("aenc.txt", "wb") as file:
            #     file.write(encrypted_data)
            with open(os.path.join(destination_dir, os.path.basename(file_path)), "wb") as file:
                file.write(encrypted_data)

            final_file = os.path.join(destination_dir, os.path.basename(file_path))
            insert_query = "INSERT INTO file_info (ownerid, findex, fhash) VALUES (%s, %s, %s)"
            cursor.execute(insert_query, (user_id, final_file, hash_value))
            conn.commit()

            # 生成并回显完整校验码
            encrypted_hash = calculate_hash(final_file)
            windowShowHash(encrypted_hash)

            status_label2.config(text="文件上传成功")


def windowShowHash(hash):
    # 创建新窗口来显示哈希值
    hash_window = tk.Toplevel(root)
    hash_window.title("这是文件完整校验码，请牢记！")
    hash_window.geometry("550x50")
    # 添加文本框显示哈希值
    hash_text = tk.Text(hash_window)
    hash_text.pack(expand=True, fill="both")
    hash_text.insert(tk.END, f"{hash}")


def showFiles():
    global files_list

    if not isUserChecked:
        status_label2.config(text="请先登录！")
        return

    user_id = user_id_entry.get()

    query = "SELECT findex FROM file_info WHERE ownerid = %s"
    cursor.execute(query, (user_id,))
    files_list = [record[0] for record in cursor.fetchall()]

    # 显示到UI
    files_display.delete(0, tk.END)
    for file_path in files_list:
        file_name = os.path.basename(file_path)
        files_display.insert(tk.END, file_name)


def downloadFile(event):
    global files_list

    if not isUserChecked:
        status_label2.config(text="请先登录！")
        return

    # 获取选定的文件索引
    selected_index = files_display.curselection()
    if selected_index:
        selected_file_index = selected_index[0]
        file_path = files_list[selected_file_index]

        if not os.path.exists(file_path):
            status_label2.config(text="文件不存在：" + file_path)
            return

        default_file_name = os.path.basename(file_path)

        # 对话框选择文件保存路径
        save_path = filedialog.asksaveasfilename(
            initialdir="/",
            title="选择保存位置",
            filetypes=(("All files", "*.*"),),
            initialfile=default_file_name
        )

        # 下载操作：先将加密文件移动到目标路径下，然后解密
        if save_path:
            try:
                shutil.copy(file_path, save_path)

                # 从数据库中查询文件的哈希值
                query = "SELECT fhash FROM file_info WHERE findex = %s"
                cursor.execute(query, (file_path,))
                hash_result = cursor.fetchone()
                # print(hash_result)
                cursor.fetchall()

                if hash_result:
                    hash_value = hash_result[0]
                    # print(hash_value)

                    key_bytes = binascii.unhexlify(hash_value)
                    key_base64 = base64.urlsafe_b64encode(key_bytes)
                    cipher_suite = Fernet(key_base64)

                    with open(save_path, "rb") as encrypted_file:
                        encrypted_data = encrypted_file.read()
                    decrypted_data = cipher_suite.decrypt(encrypted_data)
                    # print(decrypted_data)
                    with open(save_path, "wb") as decrypted_file:
                        decrypted_file.write(decrypted_data)

                status_label2.config(text="文件下载成功")
            except Exception as e:
                status_label2.config(text="文件下载失败：" + str(e))
        else:
            status_label2.config(text="未选择保存路径")
    else:
        status_label2.config(text="请选择要下载的文件")


def signupUser():
    user_id = user_id_entry.get()
    user_pwd = user_id_pw.get()

    # 检查是否存在相同的用户ID
    query = "SELECT userid FROM user_info WHERE userid = %s"
    cursor.execute(query, (user_id,))
    result = cursor.fetchone()
    if cursor.fetchone():
        print("Having selected more than one user!")
        cursor.fetchall()

    if result:
        status_label1.config(text="用户ID已存在")
    else:
        # 插入新用户数据
        insert_query = "INSERT INTO user_info (userid, userpwd) VALUES (%s, %s)"
        cursor.execute(insert_query, (user_id, user_pwd))
        conn.commit()
        status_label1.config(text="用户"+user_id+"注册成功！")

        # 创建用户专属文件夹
        user_dir = os.path.join(r"E:\A课程资料\应用密码学\EP2", user_id)
        os.makedirs(user_dir, exist_ok=True)


def checkFile():
    # 获取选定的文件索引
    selected_index = files_display.curselection()
    if selected_index:
        selected_file_index = selected_index[0]
        file_path = files_list[selected_file_index]

        if not os.path.exists(file_path):
            status_label2.config(text="文件不存在：" + file_path)
            return

        MAC = calculate_hash(file_path)
        default_file_name = os.path.basename(file_path)

        # 创建新窗口来对比哈希值
        hash_window = tk.Toplevel(root)
        hash_window.title("完整性校验："+default_file_name)
        hash_window.geometry("700x50")
        # 创建文本框用于输入用户计算的哈希值
        input_label = tk.Label(hash_window, text="输入该文件的MAC码：")
        input_label.pack(side=tk.LEFT)
        input_entry = tk.Entry(hash_window, width=70)
        input_entry.pack(side=tk.LEFT, padx=10)

        # 创建按钮用于触发比对操作
        compare_button = tk.Button(hash_window, text="比对", command=lambda: compareMAC(MAC, input_entry.get()))
        compare_button.pack(side=tk.LEFT)


def compareMAC(expected_hash, input_hash):
    print("expected_hash:", expected_hash)
    print("input_hash:", input_hash)
    if expected_hash == input_hash:
        messagebox.showinfo("完整性校验", "哈希值一致，文件完整性正常！")
        status_label2.config(text="完整性校验通过")
    else:
        messagebox.showerror("完整性校验", "哈希值不一致，文件可能已被篡改！")
        status_label2.config(text="完整性校验未通过！")


def insertKeywordHash_UI():
    # 创建新窗口
    keyword_window = tk.Toplevel(root)
    keyword_window.title("添加关键词")
    keyword_window.geometry("250x150")

    # 添加输入框和标签
    keyword1_label = tk.Label(keyword_window, text="关键词1:")
    keyword1_label.grid(row=0, column=0, padx=5, pady=5)
    keyword1_entry = tk.Entry(keyword_window)
    keyword1_entry.grid(row=0, column=1, padx=5, pady=5)

    keyword2_label = tk.Label(keyword_window, text="关键词2:")
    keyword2_label.grid(row=1, column=0, padx=5, pady=5)
    keyword2_entry = tk.Entry(keyword_window)
    keyword2_entry.grid(row=1, column=1, padx=5, pady=5)

    keyword3_label = tk.Label(keyword_window, text="关键词3:")
    keyword3_label.grid(row=2, column=0, padx=5, pady=5)
    keyword3_entry = tk.Entry(keyword_window)
    keyword3_entry.grid(row=2, column=1, padx=5, pady=5)

    # 添加确定提交按钮
    submit_button = tk.Button(keyword_window, text="确定提交", command=lambda: insertKeywordHash(keyword1_entry.get(), keyword2_entry.get(), keyword3_entry.get()))
    submit_button.grid(row=3, columnspan=2, padx=5, pady=10)


def insertKeywordHash(keyword1, keyword2, keyword3):
    # 获取选定的文件索引
    selected_index = files_display.curselection()
    if selected_index:
        selected_file_index = selected_index[0]
        selected_file_path = files_list[selected_file_index]

        # 计算关键词的哈希值
        hash1 = hashlib.sha256(keyword1.encode()).hexdigest()
        hash2 = hashlib.sha256(keyword2.encode()).hexdigest()
        hash3 = hashlib.sha256(keyword3.encode()).hexdigest()

        # 将哈希值存储到数据库中
        insert_query = "UPDATE file_info SET keyword1 = %s, keyword2 = %s, keyword3 = %s WHERE ownerid = %s AND findex = %s"
        cursor.execute(insert_query, (hash1, hash2, hash3, user_id_entry.get(), selected_file_path))
        conn.commit()

        # 提示操作成功
        messagebox.showinfo("提示", "关键词哈希值已存储！")
        status_label2.config(text="词项已更新!")
    else:
        messagebox.showerror("错误", "请选择要操作的文件")
        status_label2.config(text="请选择要操作的文件")


def searchFile_UI():
    search_window = tk.Toplevel(root)
    search_window.title("搜索文件")
    search_window.geometry("250x130")

    # 添加输入框和标签
    keyword_label = tk.Label(search_window, text="请输入关键词：")
    keyword_label.pack(pady=10)
    keyword_entry = tk.Entry(search_window, width=50)
    keyword_entry.pack(pady=10)

    # 添加搜索按钮
    search_button = tk.Button(search_window, text="搜索", command=lambda: searchFile(keyword_entry.get()))
    search_button.pack(pady=10)


def searchFile(keyword):
    if not keyword:
        messagebox.showerror("错误", "关键词不能为空")
        return

    user_id = user_id_entry.get()

    # 计算关键词的哈希值
    keyword_hash = hashlib.sha256(keyword.encode()).hexdigest()

    # 搜索数据库
    search_query = """
    SELECT findex 
    FROM file_info 
    WHERE ownerid = %s AND (keyword1 = %s OR keyword2 = %s OR keyword3 = %s)
    """
    cursor.execute(search_query, (user_id, keyword_hash, keyword_hash, keyword_hash))
    search_results = cursor.fetchall()

    # 回显搜索结果
    if search_results:
        result_files = "\n".join([os.path.basename(result[0]) for result in search_results])
        messagebox.showinfo("搜索结果", f"找到以下文件：\n{result_files}")
        status_label2.config(text="找到文件：" + result_files)
    else:
        messagebox.showinfo("搜索结果", "未找到匹配的文件")
        status_label2.config(text="没有找到文件！")






isUserChecked = False
files_list = []
# 数据库连接
conn = mysql.connector.connect(
    host="localhost",
    user="root",
    password="tang1255184",
    database="Info"
)
cursor = conn.cursor()

# 登录界
root = tk.Tk()
root.title("加密存储程序")
root.geometry("245x430")  # 设置窗口大小为650x200
# 用户ID输入框
user_id_label = tk.Label(root, text="用户ID：")
user_id_label.grid(row=0, column=0, padx=5, pady=5)
user_id_entry = tk.Entry(root)
user_id_entry.grid(row=0, column=1, columnspan=2, padx=5, pady=5)
# 用户密码输入框
user_id_pw = tk.Label(root, text="用户密码：")
user_id_pw.grid(row=1, column=0, padx=5, pady=5)
user_id_pw = tk.Entry(root, show='*')
user_id_pw.grid(row=1, column=1, columnspan=2, padx=5, pady=5)
# 登录按钮
signin_button = tk.Button(root, text="   登录   ", command=checkUser)
signin_button.grid(row=2, column=0, columnspan=1, padx=5, pady=5)
# 注册按钮
signup_button = tk.Button(root, text="   注册   ", command=signupUser)
signup_button.grid(row=2, column=1, columnspan=1, padx=5, pady=5)
# 登陆状态
status_label1 = tk.Label(root, text="尚未登录！")
status_label1.grid(row=2, column=2, columnspan=2, padx=5, pady=5)
# 水平分割线
horizontal_separator = ttk.Separator(root, orient='horizontal')
horizontal_separator.grid(row=3, columnspan=4, sticky='ew', padx=10, pady=5)
# 上传按钮
upload_button = tk.Button(root, text="上传文件", command=uploadFile)
upload_button.grid(row=4, column=0, padx=5, pady=5)
# 查看按钮
upload_button = tk.Button(root, text="查看文件", command=showFiles)
upload_button.grid(row=4, column=1, padx=5, pady=5)
# 完整性校验按钮
upload_button = tk.Button(root, text="验证文件", command=checkFile)
upload_button.grid(row=4, column=2, padx=5, pady=5)
# 词项添加按钮
upload_button = tk.Button(root, text="添加词项", command=insertKeywordHash_UI)
upload_button.grid(row=5, column=0, padx=5, pady=5)
# 可搜索加密按钮
upload_button = tk.Button(root, text="搜索文件", command=searchFile_UI)
upload_button.grid(row=5, column=2, padx=5, pady=5)
# 文件显示列表
files_display = tk.Listbox(root, width=30)
files_display.grid(row=6, column=0, columnspan=4, padx=5, pady=5)
# 状态提示
status_label2 = tk.Label(root, text="当前尚未执行文件操作")
status_label2.grid(row=7, column=0, columnspan=4, padx=5, pady=5)
# 绑定双击事件处理函数
files_display.bind("<Double-Button-1>", downloadFile)

# 窗口主循环
root.mainloop()
# 关闭游标和连
cursor.close()
conn.close()